// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/text/writer/text_writer.h>

#include <utility>
#include <pollux/common/base/pointers.h>
#include <pollux/common/encode/base64.h>
#include <pollux/exec/memory_reclaimer.h>

namespace kumo::pollux::text {
    template<typename T>
    std::optional<std::string> toTextStr(T val) {
        return std::optional(std::to_string(val));
    }

    template<>
    std::optional<std::string> toTextStr<bool>(bool val) {
        return val ? std::optional("true") : std::optional("false");
    }

    template<>
    std::optional<std::string> toTextStr<float>(float val) {
        if (std::isnan(val)) {
            return std::optional("NaN");
        } else if (std::isinf(val)) {
            return std::optional("Infinity");
        } else {
            return {std::to_string(val)};
        }
    }

    template<>
    std::optional<std::string> toTextStr<double>(double val) {
        if (std::isnan(val)) {
            return std::optional("NaN");
        } else if (std::isinf(val)) {
            return std::optional("Infinity");
        } else {
            return {std::to_string(val)};
        }
    }

    template<>
    std::optional<std::string> toTextStr<Timestamp>(Timestamp val) {
        TimestampToStringOptions options;
        options.dateTimeSeparator = ' ';
        options.precision = TimestampPrecision::kMilliseconds;
        return {val.toString(options)};
    }

    TextWriter::TextWriter(
        RowTypePtr schema,
        std::unique_ptr<dwio::common::FileSink> sink,
        const std::shared_ptr<text::WriterOptions> &options)
        : schema_(std::move(schema)),
          bufferedWriterSink_(std::make_unique<BufferedWriterSink>(
              std::move(sink),
              options->memoryPool->addLeafChild(fmt::format(
                  "{}.text_writer_node.{}",
                  options->memoryPool->name(),
                  melon::to<std::string>(melon::Random::rand64()))),
              options->defaultFlushCount)) {
    }

    void TextWriter::write(const VectorPtr &data) {
        POLLUX_CHECK_EQ(data->encoding(), VectorEncoding::Simple::ROW, "Text writer expects row vector input");
        POLLUX_CHECK(data->type()->equivalent(*schema_),
                     "The file schema type should be equal with the input row vector type.");
        const RowVector *dataRowVector = data->as<RowVector>();

        std::vector<std::shared_ptr<DecodedVector> > decodedColumnVectors;
        const auto numColumns = dataRowVector->childrenSize();
        for (size_t column = 0; column < numColumns; ++column) {
            auto decodedColumnVector = std::make_shared<DecodedVector>(DecodedVector(
                *dataRowVector->childAt(column),
                SelectivityVector(dataRowVector->size())));
            decodedColumnVectors.push_back(std::move(decodedColumnVector));
        }

        for (vector_size_t row = 0; row < data->size(); ++row) {
            for (size_t column = 0; column < numColumns; ++column) {
                if (column != 0) {
                    bufferedWriterSink_->write(TextFileTraits::kSOH);
                }
                write_cell_value(
                    decodedColumnVectors.at(column), schema_->childAt(column), row);
            }
            bufferedWriterSink_->write(TextFileTraits::kNewLine);
        }
    }

    void TextWriter::flush() {
        bufferedWriterSink_->flush();
    }

    void TextWriter::close() {
        bufferedWriterSink_->close();
    }

    void TextWriter::abort() {
        bufferedWriterSink_->abort();
    }

    void TextWriter::write_cell_value(
        const std::shared_ptr<DecodedVector> &decodedColumnVector,
        const TypePtr &type,
        vector_size_t row) {
        std::optional<std::string> dataStr;
        std::optional<StringView> dataSV;

        if (decodedColumnVector->is_null_at(row)) {
            bufferedWriterSink_->write(
                TextFileTraits::kNullData.data(), TextFileTraits::kNullData.length());
            return;
        }
        switch (type->kind()) {
            case TypeKind::BOOLEAN:
                dataStr =
                        toTextStr(melon::to<bool>(decodedColumnVector->value_at<bool>(row)));
                break;
            case TypeKind::TINYINT:
                dataStr = toTextStr(decodedColumnVector->value_at<int8_t>(row));
                break;
            case TypeKind::SMALLINT:
                dataStr = toTextStr(decodedColumnVector->value_at<int16_t>(row));
                break;
            case TypeKind::INTEGER:
                dataStr = toTextStr(decodedColumnVector->value_at<int32_t>(row));
                break;
            case TypeKind::BIGINT:
                dataStr = toTextStr(decodedColumnVector->value_at<int64_t>(row));
                break;
            case TypeKind::REAL:
                dataStr = toTextStr(decodedColumnVector->value_at<float>(row));
                break;
            case TypeKind::DOUBLE:
                dataStr = toTextStr(decodedColumnVector->value_at<double>(row));
                break;
            case TypeKind::TIMESTAMP:
                dataStr = toTextStr(decodedColumnVector->value_at<Timestamp>(row));
                break;
            case TypeKind::VARCHAR:
                dataSV = std::optional(decodedColumnVector->value_at<StringView>(row));
                break;
            case TypeKind::VARBINARY: {
                auto data = decodedColumnVector->value_at<StringView>(row);
                dataStr =
                        std::optional(encoding::Base64::encode(data.data(), data.size()));
                break;
            }
            // TODO Add support for complex types
            case TypeKind::ARRAY:
                [[fallthrough]];
            case TypeKind::MAP:
                [[fallthrough]];
            case TypeKind::ROW:
                [[fallthrough]];
            case TypeKind::UNKNOWN:
                [[fallthrough]];
            default:
                POLLUX_NYI("{} is not supported yet in TextWriter", type->kind());
        }

        if (dataStr.has_value()) {
            POLLUX_CHECK(!dataSV.has_value());
            bufferedWriterSink_->write(
                dataStr.value().data(), dataStr.value().length());
            return;
        }

        POLLUX_CHECK(dataSV.has_value());
        bufferedWriterSink_->write(dataSV.value().data(), dataSV.value().size());
    }

    std::unique_ptr<dwio::common::Writer> TextWriterFactory::createWriter(
        std::unique_ptr<dwio::common::FileSink> sink,
        const std::shared_ptr<dwio::common::WriterOptions> &options) {
        auto textOptions = std::dynamic_pointer_cast<text::WriterOptions>(options);
        POLLUX_CHECK_NOT_NULL(
            textOptions, "Text writer factory expected a Text WriterOptions object.");
        return std::make_unique<TextWriter>(
            as_row_type(options->schema), std::move(sink), textOptions);
    }

    std::unique_ptr<dwio::common::WriterOptions>
    TextWriterFactory::createWriterOptions() {
        return std::make_unique<text::WriterOptions>();
    }
} // namespace kumo::pollux::text
